#codeing=utf-8
"""
    Candidates detection:

    date: 2017.04.04
"""

from easydict import EasyDict as edict
import os
import sys
import subprocess
import time
import argparse
import pickle
import pprint

from tools.dj_test_net_vol_luna_kaggle import test_process
from tools.cdd_use_lungmask import cdd_filter_lungmask
from tools.gen_candis_list import gen_list_one_fold

from tools.nms_3d import main_nms


def parse_args():
    """
    Parse input arguments
    """
    parser = argparse.ArgumentParser(description='Detect Candidates')
    parser.add_argument('--gpu', dest='gpu_id', help='GPU id to use',
                        default=0, type=int)
    parser.add_argument('--start_idx', dest='start_scan_idx',
                        help='start_scan_idx',
                        default=None, type=int)
    parser.add_argument('--end_idx', dest='end_scan_idx',
                        help='end_scan_idx',
                        default=None, type=int)
    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(1)

    args = parser.parse_args()
    return args


if __name__ == '__main__':
    log_file_name = time.strftime('log/predict_%Y-%m-%d_%H-%M-%S.txt', time.localtime(time.time()))
    sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', 0)
    tee = subprocess.Popen(["tee", log_file_name], stdin=subprocess.PIPE)
    os.dup2(tee.stdin.fileno(), sys.stdout.fileno())
    os.dup2(tee.stdin.fileno(), sys.stderr.fileno())

    cfg = edict()

    args = parse_args()
    print('Called with args:')
    print(args)

    if args.gpu_id is not None:
        cfg.gpuid = args.gpu_id
    else:
        cfg.gpuid = 0
    if args.start_scan_idx is not None:
        cfg.start_scan_idx = args.start_scan_idx
    else:
        cfg.start_scan_idx = 0
    if args.end_scan_idx is not None:
        cfg.end_scan_idx = args.end_scan_idx
    else:
        cfg.end_scan_idx = -1

    cfg.max_iter = 60000

    cwd = os.getcwd()
    cfg.data_root_dir = cwd + '/../../../data/'
    cfg.caffemodel_path = cfg.data_root_dir + 'tmp_faster_rcnn/v15/caffemodel'
    cfg.candi_dir = cfg.data_root_dir + 'tmp_faster_rcnn/v15/candidates'
    cfg.fold_list_path = cwd + '/../fold_list_all.npy'
    cfg.kaggle_vol_path = None
    cfg.lidc_vol_path = None
    cfg.spie_vol_path = None
    cfg.stage2_vol_path = cfg.data_root_dir + 'stage2/vol.hdf5'  # config

    cfg.kaggle_lungmask_path = None
    cfg.lidc_lungmask_path = None
    cfg.spie_lungmask_path = None
    cfg.stage2_lungmask_path = cfg.data_root_dir + 'stage2/lungmask.hdf5'  # config
    cfg.candis_lungmask_dir = cfg.data_root_dir + 'tmp_faster_rcnn/v15/candidates_lungmask'
    cfg.pkl_dir = cfg.data_root_dir + 'faster_rcnn_candidates'

    # 3D NMS
    cfg.thre_xy = 5  # bbox>5mmm
    cfg.thre_dist = 10  # 10mm
    cfg.thre_score = 0.1  # this u can choose v9: 0.98 v15:0.1

    print('Using config:')
    pprint.pprint(cfg)

    fold_names = {1: 'fold1',
                  2: 'fold2',
                  3: 'fold3',
                  4: 'fold4',
                  5: 'fold5',
                  6: 'kaggle_beni',
                  7: 'kaggle_testset',
                  8: 'spie',
                  9: 'stage2'}

    print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
    print('------------------------- Testing --------------------------')
    print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
    fold_idx = 9
    fold_name = fold_names[fold_idx]
    print('----------------------- testing ' + fold_name + ' ---------------------')
    caffemodel_name = os.path.join(
        cfg.caffemodel_path, 'fold1',
        'vgg16_faster_rcnn_iter_' + str(cfg.max_iter) + '.caffemodel')
    test_process(gpu_id=cfg.gpuid,
                 caffemodel=caffemodel_name,
                 candi_dir=os.path.join(cfg.candi_dir, fold_name),
                 fold_idx=fold_idx,  # 9 for kaggle stage2 test data
                 start_scan_idx=cfg.start_scan_idx,
                 end_scan_idx=cfg.end_scan_idx,
                 fold_list_path=cfg.fold_list_path,
                 kaggle_vol_path=cfg.kaggle_vol_path,
                 lidc_vol_path=cfg.lidc_vol_path,
                 spie_vol_path=cfg.spie_vol_path,
                 stage2_vol_path=cfg.stage2_vol_path,
                 prototxt='models/pascal_voc/VGG16/faster_rcnn_end2end/test.prototxt',
                 flag_wait=True
                 )

    print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
    print('----------------------- 3D NMS -----------------------------')
    print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
    fold_idx = 9
    fold_name = fold_names[fold_idx]
    candidates_list = main_nms(candis_dir=os.path.join(cfg.candi_dir, fold_name),
                               lungmask_dir=cfg.stage2_lungmask_path,
                               vol_dir=cfg.stage2_vol_path,
                               thre_score=cfg.thre_score,
                               thre_dist=cfg.thre_dist,
                               thre_xy=cfg.thre_xy,
                               source_str=fold_name)
    if not os.path.exists(cfg.pkl_dir):
        os.makedirs(cfg.pkl_dir)
    with open(os.path.join(cfg.pkl_dir, 'candis_v15_' + fold_name + '.pkl'), 'wb') as f:
        pickle.dump(candidates_list, f, -1)
